#include "structure.h"

extern int EWINDOW;
extern int RWINDOW;
extern int BEAM;
extern double THRESH;
extern int EXTRACTOR;

extern int CL;

ostream &structure::node::write(ostream &out) const {
   out<<root<<"->"<<left<<" "<<right<<";"<<li<<","<<ri<<","<<prob<<","<<thisProb;
   return out;
}

structure::structure(int i) {
   // i leaves
   n.resize(i);
   for (int j=0; j<n.size(); ++j) {
      n[j].resize(i-j);
   }
   st.resize(i);
   en.resize(i);
}

structure::structure(const vector <string> &str, const vector <int> &s, const vector <int> &e) {
   st = s;
   en = e;
   n.resize(s.size());
   for (int j=0; j<n.size(); ++j) {
      n[j].resize(s.size()-j);
   }
   if (n.size() > 0) {
      for (int j=0; j<n[0].size(); ++j) {
	 n[0][j].push_back(node(str[j],0,0,1,1,"",""));
      }
   }
}

structure::structure(const vector <int> &s, const vector <int> &e) {
   st = s;
   en = e;
   n.resize(s.size());
   for (int j=0; j<n.size(); ++j) {
      n[j].resize(s.size()-j);
   }
}


void structure::putInternalLabel(int i, int j, const string &r, const string &ll, const string &rr) {
   assert(i < n.size());
   assert(j < n.size());
   assert(i != j);
   // internal place
   if (i < j) {
      n[j-i][i].push_back(node(r,0,0,1,1,ll,rr));
   }
   else {
      string nr=r;
      //if (ll==rr) {
	 nr += "-1";
      //}
      n[i-j][j].push_back(node(nr,0,0,1,1,rr,ll));
   }
}
 
void structure::putBottomLabel(int i, const string &r, int s, int e) {
   assert(i < n.size());
   // leaf place
   n[0][i].push_back(node(r,0,0,1,1,"",""));
   st[i] = s;
   en[i] = e;
}

void structure::fillIn() {
   for (int i=1; i<n.size(); ++i) {
      for (int j=0; j<n[i].size(); ++j) {
	 if (n[i][j].size()==0) {
	    string r = "NR_" + n[0][j][0].root + "_" + n[0][i+j][0].root;
	    n[i][j].push_back(node(r,0,0,1,1,n[0][j][0].root,n[0][i+j][0].root));
	 }
      }
   }
}

void structure::giveFirstRelEntLabels(vector <string> &pr, vector <int> &ps, vector <int> &pe) const {
   if (n.size() > 0) {
      giveFirstRelEntLabels_helper(pr,ps,pe,n.size()-1,0,0); 
   }
}

bool structure::entityMatch(int s, int e, const string &r) const {
   int i;
   for (i=0; i<st.size(); ++i) {
      if (st[i]==s && en[i]==e) {
	 break;
      }
   }
   if (i == st.size()) return false;
   return (n[0][i][0].root==r);
}

void structure::giveFirstRelEntLabels_helper(vector <string> &pr, vector <int> &ps, vector <int> &pe, int i, int j, int k) const {
   if (i > 0) {
      if (n[i][j].size() > k) {
         if (n[i][j][k].root[0]=='N' && n[i][j][k].root[1]=='R') {
	 }
         else {
	    string r=n[i][j][k].left;
	    int s=st[j], e=en[j];
	    bool f=true;
	    for (int a=0; a<pr.size(); ++a) {
	       if (s==ps[a] && e==pe[a] && r==pr[a]) {
		  f=false;
		  break;
	       }
	    }
	    if (f) {
	       ps.push_back(s);
	       pe.push_back(e);
	       pr.push_back(r);
	    }
	    r=n[i][j][k].right;
	    s=st[i+j], e=en[i+j];
	    f=true;
	    for (int a=0; a<pr.size(); ++a) {
	       if (s==ps[a] && e==pe[a] && r==pr[a]) {
		  f=false;
		  break;
	       }
	    }
	    if (f) {
	       ps.push_back(s);
	       pe.push_back(e);
	       pr.push_back(r);
	    }

	 }
	 giveFirstRelEntLabels_helper(pr,ps,pe,i-1,j,n[i][j][k].li);
	 giveFirstRelEntLabels_helper(pr,ps,pe,i-1,j+1,n[i][j][k].ri);
      }
      else {
	 giveFirstRelEntLabels_helper(pr,ps,pe,i-1,j,0);
	 giveFirstRelEntLabels_helper(pr,ps,pe,i-1,j+1,0);
      }
   }
}


void structure::giveFirstLabels_helper(vector <string> &pr, vector <int> &ps1, vector <int> &pe1, vector <int> &ps2, vector <int> &pe2, int i, int j, int k)  const {
   if (n[i][j].size() > k) {
      string r=n[i][j][k].root;
      int s1=st[j],e1=en[j],s2=st[i+j],e2=en[i+j];

      bool f=true;
      for (int a=0; a<pr.size(); ++a) {
	 if (s1==ps1[a] && e1==pe1[a] && s2==ps2[a] && e2==pe2[a] && r==pr[a]) {
	    f=false;
	    break;
	 }
      }
      if (f) {
         pr.push_back(r);
         ps1.push_back(s1);
         pe1.push_back(e1);
         ps2.push_back(s2);
         pe2.push_back(e2);
      }
      if (i > 0) {
	 giveFirstLabels_helper(pr,ps1,pe1,ps2,pe2,i-1,j,n[i][j][k].li);
	 giveFirstLabels_helper(pr,ps1,pe1,ps2,pe2,i-1,j+1,n[i][j][k].ri);
      }
   }
   else {
      giveFirstLabels_helper(pr,ps1,pe1,ps2,pe2,i-1,j,0);
      giveFirstLabels_helper(pr,ps1,pe1,ps2,pe2,i-1,j+1,0);
   }
}

void structure::giveFirstLabels(vector <string> &pr, vector <int> &ps1, vector <int> &pe1, vector <int> &ps2, vector <int> &pe2) const {
   if (n.size() > 0) {
      giveFirstLabels_helper(pr,ps1,pe1,ps2,pe2,n.size()-1,0,0); 
   }
}

void structure::giveLeafLabels(vector <int> &s, vector <int> &e, vector <int> &r, int i, int j, int k) const {
   if (i < 0) return;
   if (i==0) {
      if (n[i][j].size() > k) {
      int ss=st[j];
      int ee=en[j];
      int rr = mapInsert(n[i][j][k].root);

      bool f=true;
      for (int a=0; a<s.size(); ++a) {
	 if (s[a]==ss && e[a]==ee && r[a]==rr) {
	    f=false;
	    break;
	 }
      }

      if (f) {
	 s.push_back(ss);
	 e.push_back(ee);
	 r.push_back(rr);
      }
      }
   }
   else {
      if (n[i][j].size() > k) {
	 giveLeafLabels(s,e,r,i-1,j,n[i][j][k].li);
	 giveLeafLabels(s,e,r,i-1,j+1,n[i][j][k].ri);
      }
      else {
	 giveLeafLabels(s,e,r,i-1,j,0);
	 giveLeafLabels(s,e,r,i-1,j+1,0);
      }
   }
}

void structure::giveBottomEgs(vector <vector <string> > &pr, vector <int> &ps, vector <int> &pe) const {
   if (n.size() > 0) {
      for (int j=0; j<n[0].size(); ++j) {
	 for (int k=0; k<n[0][j].size(); ++k) {
	    vector <string> s;
	    s.push_back(n[0][j][k].root);
	    s.push_back("entity");
	    pr.push_back(s);
	    ps.push_back(st[j]);
	    pe.push_back(en[j]);
	 }
      }
   }
}
void structure::giveInternalEgs(vector <vector <string> > &pr, vector <vector <int> > &ps, vector <vector <int> > &pe, vector <vector <vector <string> > > &path, vector <vector <vector <double> > > &pathProb) const {
   for (int i=1; i<n.size(); ++i) {
      for (int j=0; j<n[i].size(); ++j) {
	 for (int k=0; k<n[i][j].size(); ++k) {
	    vector <string> s;
	    s.push_back(n[i][j][k].root);
	    s.push_back(n[i][j][k].left);
	    s.push_back(n[i][j][k].right);
	    pr.push_back(s);
	    vector <vector <string> > p;
	    vector <vector <double> > pp;
	    givePathInfo(i,j,k,p,pp);
	    path.push_back(p);
	    pathProb.push_back(pp);
	    vector <int> ss,ee;
	    for (int a=j; a<=j+i; ++a) {
	       ss.push_back(st[a]);
	       ee.push_back(en[a]);
	    }
	    ps.push_back(ss);
	    pe.push_back(ee);
	 }
      }
   }
}

void structure::givePathInfo(int i, int j, int k, vector <vector <string> > &path, vector <vector <double> > &pathProb) const {
   if (i > 1) {
      for (int a=j+1; a<=j+i-1; ++a) {
	 vector <string> ll;
	 vector <double> pp;
	 string l1;
	 double p1;
	 giveNodeInfo(i,j,k,j,a,l1,p1);
	 ll.push_back(l1);
	 pp.push_back(p1);
	 string l2;
	 double p2;
	 giveNodeInfo(i,j,k,a,j+i,l2,p2);
	 ll.push_back(l2);
	 pp.push_back(p2);

	 //if (find(path.begin(),path.end(),ll)==path.end()) {
	    path.push_back(ll);
	    pathProb.push_back(pp);
	 //}
      }
   }
}

void structure::giveNodeInfo(int i, int j, int k, int s, int e, string &l, double &tp) const {
   // gives the label and probability of the node from s to e under n[i][j][k] node
   assert(s < e);
   assert(e <= i+j);
   assert(s >= j);

   if (s==j) {
      if (e==i+j) {
	 l = n[i][j][k].root;
	 tp = n[i][j][k].thisProb;
	 return;
      }
      else { // search on the left side
	 giveNodeInfo(i-1,j,n[i][j][k].li,s,e,l,tp);
      }
   }
   else {
      // search on the right side
      giveNodeInfo(i-1,j+1,n[i][j][k].ri,s,e,l,tp);
   }
}

void structure::giveProds(vector <vector <string> > &pr) const {
   for (int i=1; i<n.size(); ++i) {
      for (int j=0; j<n[i].size(); ++j) {
	 for (int k=0; k<n[i][j].size(); ++k) {
	    vector <string> s;
	    s.push_back(n[i][j][k].root);
	    s.push_back(n[i][j][k].left);
	    s.push_back(n[i][j][k].right);
	    pr.push_back(s);
	 }
      }
   }
   if (n.size() > 0) {
      for (int j=0; j<n[0].size(); ++j) {
         for (int k=0; k<n[0][j].size(); ++k) {
	    vector <string> s;
	    s.push_back(n[0][j][k].root);
	    s.push_back("entity");
	    pr.push_back(s);
         }
      }
   }
}


void structure::parse(vector <classifier> &cl, vector <int> &words, vector <int> &postags) {
   pn tsw; // to store examples encountered in this parsing
   // first the entities
   if (n.size() > 0 && EXTRACTOR != RELATION_ONLY) {
      assert(st.size()==n[0].size());
      assert(en.size()==n[0].size());
      for (int j=0; j<n[0].size(); ++j) {
	 if (n[0][j].size() > 0) continue;
	 vector <int> w,pt,bw,aw,bp,ap,ntw,ntp;
	 for (int k=st[j]; k<=en[j]; ++k) {
	    w.push_back(words[k]);
	    pt.push_back(postags[k]);
	 }
	 int b=st[j]-EWINDOW;
	 if (b < 0) b=0;
	 for (int k=b; k<st[j]; ++k) {
	    bw.push_back(words[k]);
	    bp.push_back(postags[k]);
	 }
	 int a=en[j]+EWINDOW;
	 if (a > (int)words.size()-1) a = (int)words.size()-1; 
	 for (int k=en[j]+1; k<=a; ++k) {
	    aw.push_back(words[k]);
	    ap.push_back(postags[k]);
	 }
	 //ntw=w;
	 //ntp=pt;
	 int ind = -1;

	 for (int c=0; c<cl.size(); ++c) {
	    if (cl[c].isTerm()) {
	       if (ind==-1) {
		  posNeg eg(w,pt,bw,aw,bp,ap,ntw,ntp,true);
		  ind = tsw.add(eg);
	       }
	       //if (cl[c].giveDec(ind,tsw) > 0) {
	       double prob;
	       //if (cl[c].giveLhs()=="Other") prob=0.5;
	       //else 
		  prob = cl[c].giveP(ind,tsw);
	          //if (prob > THRESH) 
		     n[0][j].push_back(node(cl[c].giveLhs(),0,0,prob,prob,"",""));
	       //}
	    }
	 }
	 sort(n[0][j].begin(),n[0][j].end());
	 if (n[0][j].size() > BEAM) n[0][j].resize(BEAM);
	 if (n[0][j].size()==0) {
	    n[0][j].push_back(node("Other",0,0,1,1,"",""));
	 }
      }
   }
   if (EXTRACTOR != ENTITY_ONLY) {
      for (int i=1; i<n.size(); ++i) {
         for (int j=0; j<n[i].size(); ++j) {
	    if (n[i][j].size() > 0) continue;
	    vector <int> w,pt,bw,aw,bp,ap,ntw,ntp;
	    vector <int> rl;
	    vector <int> sl,el;
	    giveLeafLabels(sl,el,rl,i,j,0);

	    for (int k=st[j]; k<=en[i+j]; ++k) {
	       w.push_back(words[k]);
	       pt.push_back(postags[k]);
	    }
	    int b=st[j]-RWINDOW;
	    if (b < 0) b=0;
	    for (int k=b; k<st[j]; ++k) {
	       bw.push_back(words[k]);
	       bp.push_back(postags[k]);
	    }
	    int a=en[i+j]+RWINDOW;
	    if (a > (int)words.size()-1) a = (int)words.size()-1; 
	    for (int k=en[i+j]+1; k<=a; ++k) {
	       aw.push_back(words[k]);
	       ap.push_back(postags[k]);
	    }
	    //ntw.push_back(0); // will be relplaced
	    //ntp.push_back(0);
	    for (int k=en[j]+1; k<st[i+j]; ++k) {
	       ntw.push_back(words[k]);
	       ntp.push_back(postags[k]);
	    }
	    //ntw.push_back(0); // will be replaced
	    //ntp.push_back(0);

	    for (int f=0; f<n[i-1][j].size(); ++f) {
	       for (int g=0; g<n[i-1][j+1].size(); ++g) {
		  if (i > 1) {
		     if (n[i-1][j][f].ri != n[i-1][j+1][g].li) continue; // inconsistent overlap
		  }
		  vector <string> rhs;
		  if (i > 1) {
		     rhs.push_back(n[i-1][j][f].left);
		     rhs.push_back(n[i-1][j+1][g].right);
		  }
		  else {
		     rhs.push_back(n[0][j][f].root);
		     rhs.push_back(n[0][j+1][g].root);
		  }
		  int ind = -1;
	          for (int c=0; c<cl.size(); ++c) {
		     if (cl[c].equalsRhs(rhs)) {
			//ntw[0] = mapInsert(rhs[0]);
			//ntp[0] = ntw[0];
			//ntw[ntw.size()-1] = mapInsert(rhs[1]);
			//ntp[ntp.size()-1] = ntw[ntw.size()-1];
			if (ind == -1) {
			   vector <int> w,pt,bw,aw,bp,ap,ntw,ntp;
			   vector <int> rl;
			   vector <int> sl,el;
			   giveLeafLabels(sl,el,rl,i-1,j,f);
			   giveLeafLabels(sl,el,rl,i-1,j+1,g);
	    
			   for (int k=st[j]; k<=en[i+j]; ++k) {
			      int a;
			      for (a=0; a<sl.size(); ++a) {
				 if (k==sl[a]) {
				    if ((sl[a]==st[j] && el[a]==en[j])||
					  (sl[a]==st[i+j] && el[a]==en[i+j])) 
				    { }
				    else break;
				 }
			      }
			      if (a < sl.size()) {
				 ntw.push_back(rl[a]);
				 ntp.push_back(rl[a]);
				 k = el[a];
			      }
			      else {
				 w.push_back(words[k]);
				 pt.push_back(postags[k]);
			      }
			   }
	                   int b=st[j]-RWINDOW;
			   if (b < 0) b=0;
			   for (int k=b; k<st[j]; ++k) {
			      bw.push_back(words[k]);
			      bp.push_back(postags[k]);
			   }
	   
			   int a=en[i+j]+RWINDOW;
			   if (a > (int)words.size()-1) a = (int)words.size()-1; 
			   for (int k=en[i+j]+1; k<=a; ++k) {
			      aw.push_back(words[k]);
			      ap.push_back(postags[k]);
			   }
			   //ntw.push_back(0); // will be relplaced
			   //ntp.push_back(0);
			   for (int k=en[j]+1; k<st[i+j]; ++k) {
			      int a;
			      for (a=0; a<sl.size(); ++a) {
				 if (k==sl[a]) {
				    break;
				 }
			      }
			      if (a < sl.size()) {
				 ntw.push_back(rl[a]);
				 ntp.push_back(rl[a]);
				 k = el[a];
			      }
			      else {
				 ntw.push_back(words[k]);
				 ntp.push_back(postags[k]);
			      }
			   }
			   //ntw.push_back(0); // will be replaced
			   //ntp.push_back(0);
			   posNeg eg(w,pt,bw,aw,bp,ap,ntw,ntp,false);

			   vector <vector <string> > path;
			   vector <vector <double> > pathProb;

			   for (int a=j+1; a<=j+i-1; ++a) {
			      vector <string> pa(2);
			      vector <double> pp(2);
			      giveNodeInfo(i-1,j,f,j,a,pa[0],pp[0]);
			      giveNodeInfo(i-1,j+1,g,a,j+i,pa[1],pp[1]);
			      //if (find(path.begin(),path.end(),pa)==path.end()) {
			         path.push_back(pa);
			         pathProb.push_back(pp);
			      //}
			   }

			   eg.addPath(path,pathProb);
			      
			   ind = tsw.add(eg);
			}
			//if (cl[c].giveDec(ind,tsw) > 0) {
			   double prob, tprob;
			   if (cl[c].ifNr()) {
			      //prob = cl[c].giveP(ind,tsw);
			      //prob = cl[CL].giveP(ind,tsw);
			      prob = 0.5;
			   }
			   else {
			      prob = cl[c].giveP(ind,tsw);
			      //if (prob < 0.1) prob=0;
			   }
			   //if (prob > THRESH) {
			   tprob=prob;
			      prob *= n[i-1][j][f].prob;
			      prob *= n[i-1][j+1][g].prob;
			      if (i > 1) {
				 double a=n[i-2][j+1][n[i-1][j][f].ri].prob;
				 if (a != 0) prob /= a;
			      }
			      //if (prob > THRESH*0.1) 
			         n[i][j].push_back(node(cl[c].giveLhs(),f,g,prob,tprob,rhs[0],rhs[1]));
			//}
		     }
		  }
	       }
	    }
	    sort(n[i][j].begin(),n[i][j].end());
	    if (n[i][j].size() > BEAM) n[i][j].resize(BEAM);
	    if (n[i][j].size()==0) {
	       n[i][j].push_back(node("NR",0,0,1,1,"",""));
	    }
	 }
      }
   }
   else {
      for (int i=1; i<n.size(); ++i) {
         for (int j=0; j<n[i].size(); ++j) {
	    n[i][j].push_back(node("NR",0,0,1,1,"",""));
	 }
     }
   }
}


ostream &structure::write(ostream &out) const {
   for (int i=0; i<n.size(); ++i) {
      for (int j=0; j<n[i].size(); ++j) {
	 for (int k=0; k<n[i][j].size(); ++k) {
	    out<<"(";
	    n[i][j][k].write(out);

	    vector <vector <string> > path;
	    vector <vector <double> > pathProb;
	    givePathInfo(i,j,k,path,pathProb);

	    for (int a=0; a<path.size(); ++a) {
	       out<<"["<<path[a][0]<<" "<<pathProb[a][0]<<" - ";
	       out<<path[a][1]<<" "<<pathProb[a][1]<<"] ";
	    }

	    out<<") ";
	 }
	 out<<"\n";
      }
      out<<"\n\n";
   }
   return out;
}
